import torch
from torch import nn
from torchvision.datasets import ImageFolder  # 直接加载数据目录，进行分类
from torchvision.transforms import transforms  # 将图像进行一定处理的工具
from torch.utils.data import DataLoader

# ==================== 图像数据集处理 ===================
# 对单张图片进行处理
transform = transforms.Compose([
    transforms.ToTensor()  # 将加载的图像数据转化为pytorch.tensor类型
])

# ImageFoler将指定目录下的文件数据进行自动分类
dataset = ImageFolder("./data/numbers", transform=transform)

# 将上述数据集进行 将数据分批次、将数据进行随机打散
dataloader = DataLoader(dataset, batch_size=1000, shuffle=True)

# ==================== 搭建训练模型（卷积神经网络） ===================
# 训练手写数字数据集 （28，28，3） --> (6-8,6-8,n) 最终卷积之后的图像大小必须在 6-8 之间（重点）
model = nn.Sequential(
    nn.Conv2d(3, 16, 3, 1, 1),  # (28x28x16)
    nn.ReLU(),
    nn.MaxPool2d(2, 2),  # (14x14x16)
    nn.Conv2d(16, 32, 3, 1, 1),  # (14x14x32)
    nn.ReLU(),
    nn.MaxPool2d(2, 2),  # (7x7x32)
    nn.Flatten(),  # (1568,) 将上述的图像矩阵进行拉伸（展平）
    nn.Dropout(),  # 随机削减50% 的神经元数量（训练提速）
    nn.Linear(7 * 7 * 32, 1024),
    nn.ReLU(),
    nn.Linear(1024, 10),  # 10是分类的数量
    nn.LogSoftmax(dim=-1)  # 多分类激活函数
)

# ==================== 开始训练 ===================
# 损失函数
criterion = nn.CrossEntropyLoss()
# 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# 模型训练次数
epochs = 10000
# 开始训练
for epoch in range(epochs):
    losses = [] # 记录一个epoch下每一个batch_size的损失
    # 根据batch_size训练数据
    for features, labels in dataloader:
        optimizer.zero_grad()
        # 预测
        label_predict = model(features)
        # 计算损失
        loss = criterion(label_predict, labels.long())  # 此处labels需要变为整数
        loss.backward()  # 反向传播
        losses.append(loss.item())
        # 更新w、b的参数
        optimizer.step()

    # 计算平均损失
    avg_loss = sum(losses) / len(losses)
    print(f"{epoch + 1} / {epochs} -- loss:{avg_loss:.4f}")  # :.4f保留四位小数

# ==================== 保存模型 ===================
# 将训练好的模型参数进行保存
torch.save(model.state_dict(),"./save/best.pt")